# 0. 引入必要的包
# TODO
import os
import numpy as np
import cv2 as cv
from sklearn.model_selection import train_test_split
from util import get, preprocess_image, dump, load
from rich.progress import track
import time

# 1. 读取配置文件中的信息
train_dir = get("train") # 获取 训练数据路径
char_styles = get("char_styles") # 获取 字符样式列表，注意: 必须是列表
new_size = get("new_size") # 获取 新图像大小元组, 注意: 必须包含h和w
Xy_root = get('Xy_root') # 获取训练集的路径

# 2. 生成X,y
print("# 读取训练数据并进行预处理，")
img_name_list = os.listdir(train_dir)
img_list = []
label_list = []
for name, _ in zip(img_name_list, track(range(len(img_name_list)), description="图片处理进度")):
    img_list.append(preprocess_image(f'{train_dir}/{name}', new_size=new_size))
    label_list.append([s for s in char_styles if s in name][0])
    pass
print("构建训练集X及其标签y...")
X = np.concatenate(img_list, axis=0).reshape(len(img_list), -1)
y = np.array([char_styles.index(s) for s in label_list]).reshape(-1, 1)
print("构建训练集X及其标签y成功！")
# TODO

# 3. 分割测试集和训练集
print("# 将数据按 80% 和 20% 的比例分割")
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

# 4. 打印样本维度和类型信息
print("X_train: ", X_train.shape, X_train.dtype)  # 训练集特征的维度和类型
print("X_test: ", X_test.shape, X_test.dtype)  # 测试集特征的维度和类型
print("y_train: ", y_train.shape, y_train.dtype)  # 训练集标签的维度和类型
print("y_test: ", y_test.shape, y_test.dtype)  # 测试集标签的维度和类型

# 5. 序列化分割后的训练和测试样本
# TODO
dump((X_train, X_test, y_train, y_test), 'Xy', f'{Xy_root}')